{"nbformat":4,"nbformat_minor":0,"metadata":{"accelerator":"GPU","colab":{"name":"“rnn_lstm_attention_captioning.ipynb”的副本","provenance":[{"file_id":"1BK802C-gxfbMs_udm6WRDoH9Lkx5RxTd","timestamp":1606219642170}],"collapsed_sections":[],"toc_visible":true},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.7.3"}},"cells":[{"cell_type":"markdown","metadata":{"id":"DDJwQPZcupab"},"source":["# EECS 498-007/598-005 Assignment 4-2: RNN, LSTM, and Attention for Image Captioning\n","\n","Before we start, please put your name and UMID in following format\n","\n",": Firstname LASTNAME, #00000000   //   e.g.) Justin JOHNSON, #12345678"]},{"cell_type":"markdown","metadata":{"id":"2KMxqLt1h2kx"},"source":["**Your Answer:**   \n","Hello WORLD, #XXXXXXXX"]},{"cell_type":"markdown","metadata":{"id":"Euek3FWn6bhA","tags":["pdf-title"]},"source":["# Image Captioning with RNNs\n","In this exercise you will implement a vanilla recurrent neural networks and use them it to train a model that can generate novel captions for images."]},{"cell_type":"markdown","metadata":{"id":"ubB_0e-UAOVK"},"source":["## Install starter code\n","We will continue using the utility functions that we've used for previous assignments: [`coutils` package](https://github.com/deepvision-class/starter-code). Run this cell to download and install it.\n"]},{"cell_type":"code","metadata":{"id":"ASkY27ZtA7Is"},"source":["!pip install git+https://github.com/deepvision-class/starter-code"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"MzqbYcKdz6ew"},"source":["## Setup code\n","Run some setup code for this notebook: Import some useful packages and increase the default figure size."]},{"cell_type":"code","metadata":{"id":"q53DlMXboP-T"},"source":["import math\n","import torch\n","from torch import nn\n","from torch.nn.parameter import Parameter\n","import torch.nn.functional as F\n","import coutils\n","from coutils import fix_random_seed, rel_error, compute_numeric_gradient, \\\n","                    tensor_to_image, decode_captions, attention_visualizer\n","import matplotlib.pyplot as plt\n","import time\n","\n","# for plotting\n","%matplotlib inline\n","plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots\n","plt.rcParams['image.interpolation'] = 'nearest'\n","plt.rcParams['image.cmap'] = 'gray'\n","\n","# data type and device for torch.tensor\n","to_float = {'dtype': torch.float, 'device': 'cpu'}\n","to_float_cuda = {'dtype': torch.float, 'device': 'cuda'}\n","to_double = {'dtype': torch.double, 'device': 'cpu'}\n","to_double_cuda = {'dtype': torch.double, 'device': 'cuda'}\n","to_long = {'dtype': torch.long, 'device': 'cpu'}\n","to_long_cuda = {'dtype': torch.long, 'device': 'cuda'}"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"OvUDZWGU3VLV"},"source":["We will use GPUs to accelerate our computation in this notebook. Run the following to make sure GPUs are enabled:"]},{"cell_type":"code","metadata":{"id":"RrAX9FOLpr9k"},"source":["if torch.cuda.is_available:\n","  print('Good to go!')\n","else:\n","  print('Please set GPU via Edit -> Notebook Settings.')"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"WCPZwvOd6bhF","tags":["pdf-ignore"]},"source":["# Microsoft COCO\n","For this exercise we will use the 2014 release of the [Microsoft COCO dataset](http://mscoco.org/) which has become the standard testbed for image captioning. The dataset consists of 80,000 training images and 40,000 validation images, each annotated with 5 captions written by workers on Amazon Mechanical Turk.\n","\n","We have preprocessed the data for you already and saved them into a serialized data file. It contains 10,000 image-caption pairs for training and 500 for testing. The images have been downsampled to 112x112 for computation efficiency and captions are tokenized and numericalized, clamped to 15 words. You can download the file named `coco.pt` (378MB) with the link below and run some useful stats.\n","\n","You will later use MobileNet v2 to extract features for the images. A few notes on the caption preprocessing:\n","\n","Dealing with strings is inefficient, so we will work with an encoded version of the captions. Each word is assigned an integer ID, allowing us to represent a caption by a sequence of integers. The mapping between integer IDs and words is saved in an entry named `vocab` (both `idx_to_token` and `token_to_idx`), and we use the function `decode_captions` from the starter code to convert tensors of integer IDs back into strings.\n","\n","There are a couple special tokens that we add to the vocabulary. We prepend a special `<START>` token and append an `<END>` token to the beginning and end of each caption respectively. Rare words are replaced with a special `<UNK>` token (for \"unknown\"). In addition, since we want to train with minibatches containing captions of different lengths, we pad short captions with a special `<NULL>` token after the `<END>` token and don't compute loss or gradient for `<NULL>` tokens. Since they are a bit of a pain, we have taken care of all implementation details around special tokens for you."]},{"cell_type":"code","metadata":{"id":"IMok4gFXqjre"},"source":["# Download and load serialized COCO data from coco.pt\n","# It contains a dictionary of\n","# \"train_images\" - resized training images (112x112)\n","# \"val_images\" - resized validation images (112x112)\n","# \"train_captions\" - tokenized and numericalized training captions\n","# \"val_captions\" - tokenized and numericalized validation captions\n","# \"vocab\" - caption vocabulary, including \"idx_to_token\" and \"token_to_idx\"\n","\n","!wget http://web.eecs.umich.edu/~justincj/teaching/eecs498/coco.pt\n","data_dict = torch.load('coco.pt')\n","\n","# print out all the keys and values from the data dictionary\n","for k, v in data_dict.items():\n","    if type(v) == torch.Tensor:\n","        print(k, type(v), v.shape, v.dtype)\n","    else:\n","        print(k, type(v), v.keys())\n","\n","num_train = data_dict['train_images'].size(0)\n","num_val = data_dict['val_images'].size(0)\n","assert data_dict['train_images'].size(0) == data_dict['train_captions'].size(0) and \\\n","       data_dict['val_images'].size(0) == data_dict['val_captions'].size(0), \\\n","       'shapes of data mismatch!'\n","\n","print('\\nTrain images shape: ', data_dict['train_images'].shape)\n","print('Train caption tokens shape: ', data_dict['train_captions'].shape)\n","print('Validation images shape: ', data_dict['val_images'].shape)\n","print('Validation caption tokens shape: ', data_dict['val_captions'].shape)\n","print('total number of caption tokens: ', len(data_dict['vocab']['idx_to_token']))\n","print('mappings (list) from index to caption token: ', data_dict['vocab']['idx_to_token'])\n","print('mappings (dict) from caption token to index: ', data_dict['vocab']['token_to_idx'])\n","\n","# declare variables for special tokens\n","NULL_index = data_dict['vocab']['token_to_idx']['<NULL>']\n","START_index = data_dict['vocab']['token_to_idx']['<START>']\n","END_index = data_dict['vocab']['token_to_idx']['<END>']\n","UNK_index = data_dict['vocab']['token_to_idx']['<UNK>']"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"80RW_nSH6bhH"},"source":["## Look at the data\n","It is always a good idea to look at examples from the dataset before working with it.\n","\n","Run the following to sample a small minibatch of training data and show the images and their captions. Running it multiple times and looking at the results helps you to get a sense of the dataset.\n","\n","Note that we decode the captions using the `decode_captions` function."]},{"cell_type":"code","metadata":{"id":"l-oiW9Ut6bhH"},"source":["# Sample a minibatch and show the reshaped 112x112 images and captions\n","batch_size = 3\n","\n","sample_idx = torch.randint(0, num_train, (batch_size,))\n","sample_images = data_dict['train_images'][sample_idx]\n","sample_captions = data_dict['train_captions'][sample_idx]\n","for i in range(batch_size):\n","  plt.imshow(sample_images[i].permute(1, 2, 0))\n","  plt.axis('off')\n","  caption_str = decode_captions(sample_captions[i], data_dict['vocab']['idx_to_token'])\n","  plt.title(caption_str)\n","  plt.show()"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"b2SQMNIH6bhJ"},"source":["# Recurrent Neural Networks\n","As discussed in lecture, we will use Recurrent Neural Network (RNN) language models for image captioning. We will cover the vanilla RNN model first and later LSTM and attention-based language models."]},{"cell_type":"markdown","metadata":{"id":"6XHZMI356bhJ"},"source":["## Vanilla RNN: step forward\n","First implement the forward pass for a single timestep of a vanilla recurrent neural network."]},{"cell_type":"code","metadata":{"id":"NoRPdqKihJjy"},"source":["def rnn_step_forward(x, prev_h, Wx, Wh, b):\n","    \"\"\"\n","    Run the forward pass for a single timestep of a vanilla RNN that uses a tanh\n","    activation function.\n","\n","    The input data has dimension D, the hidden state has dimension H, and we use\n","    a minibatch size of N.\n","\n","    Inputs:\n","    - x: Input data for this timestep, of shape (N, D).\n","    - prev_h: Hidden state from previous timestep, of shape (N, H)\n","    - Wx: Weight matrix for input-to-hidden connections, of shape (D, H)\n","    - Wh: Weight matrix for hidden-to-hidden connections, of shape (H, H)\n","    - b: Biases, of shape (H,)\n","\n","    Returns a tuple of:\n","    - next_h: Next hidden state, of shape (N, H)\n","    - cache: Tuple of values needed for the backward pass.\n","    \"\"\"\n","    next_h, cache = None, None\n","    ##############################################################################\n","    # TODO: Implement a single forward step for the vanilla RNN. Store the next  #\n","    # hidden state and any values you need for the backward pass in the next_h   #\n","    # and cache variables respectively.                                          #\n","    ##############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","    return next_h, cache"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"TLhpgKhlhSse"},"source":["Run the following to check your implementation. You should see errors on the order of 1e-8 or less."]},{"cell_type":"code","metadata":{"id":"c3oU8JJj6bhK"},"source":["N, D, H = 3, 10, 4\n","\n","x = torch.linspace(-0.4, 0.7, steps=N*D, **to_double_cuda).reshape(N, D)\n","prev_h = torch.linspace(-0.2, 0.5, steps=N*H, **to_double_cuda).reshape(N, H)\n","Wx = torch.linspace(-0.1, 0.9, steps=D*H, **to_double_cuda).reshape(D, H)\n","Wh = torch.linspace(-0.3, 0.7, steps=H*H, **to_double_cuda).reshape(H, H)\n","b = torch.linspace(-0.2, 0.4, steps=H, **to_double_cuda)\n","\n","next_h, _ = rnn_step_forward(x, prev_h, Wx, Wh, b)\n","expected_next_h = torch.tensor([\n","  [-0.58172089, -0.50182032, -0.41232771, -0.31410098],\n","  [ 0.66854692,  0.79562378,  0.87755553,  0.92795967],\n","  [ 0.97934501,  0.99144213,  0.99646691,  0.99854353]], **to_double_cuda)\n","\n","print('next_h error: ', rel_error(expected_next_h, next_h))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"tid-ljPA6bhL"},"source":["## Vanilla RNN: step backward\n","Then implement the backward pass for a single timestep of a vanilla recurrent neural network."]},{"cell_type":"code","metadata":{"id":"_VpwjpiAhdqm"},"source":["def rnn_step_backward(dnext_h, cache):\n","    \"\"\"\n","    Backward pass for a single timestep of a vanilla RNN.\n","\n","    Inputs:\n","    - dnext_h: Gradient of loss with respect to next hidden state, of shape (N, H)\n","    - cache: Cache object from the forward pass\n","\n","    Returns a tuple of:\n","    - dx: Gradients of input data, of shape (N, D)\n","    - dprev_h: Gradients of previous hidden state, of shape (N, H)\n","    - dWx: Gradients of input-to-hidden weights, of shape (D, H)\n","    - dWh: Gradients of hidden-to-hidden weights, of shape (H, H)\n","    - db: Gradients of bias vector, of shape (H,)\n","    \"\"\"\n","    dx, dprev_h, dWx, dWh, db = None, None, None, None, None\n","    ##############################################################################\n","    # TODO: Implement the backward pass for a single step of a vanilla RNN.      #\n","    #                                                                            #\n","    # HINT: For the tanh function, you can compute the local derivative in terms #\n","    # of the output value from tanh.                                             #\n","    ##############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","    return dx, dprev_h, dWx, dWh, db"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"K0ALCpQMhpKP"},"source":["Run the following to numerically gradient check your implementation. You should see errors on the order of `1e-8` or less."]},{"cell_type":"code","metadata":{"id":"KPyfJofC6bhM"},"source":["fix_random_seed(0)\n","N, D, H = 4, 5, 6\n","x = torch.randn(N, D, **to_double_cuda)\n","h = torch.randn(N, H, **to_double_cuda)\n","Wx = torch.randn(D, H, **to_double_cuda)\n","Wh = torch.randn(H, H, **to_double_cuda)\n","b = torch.randn(H, **to_double_cuda)\n","\n","out, cache = rnn_step_forward(x, h, Wx, Wh, b)\n","\n","dnext_h = torch.randn(*out.shape, **to_double_cuda)\n","\n","fx = lambda x: rnn_step_forward(x, h, Wx, Wh, b)[0]\n","fh = lambda h: rnn_step_forward(x, h, Wx, Wh, b)[0]\n","fWx = lambda Wx: rnn_step_forward(x, h, Wx, Wh, b)[0]\n","fWh = lambda Wh: rnn_step_forward(x, h, Wx, Wh, b)[0]\n","fb = lambda b: rnn_step_forward(x, h, Wx, Wh, b)[0]\n","\n","dx_num = compute_numeric_gradient(fx, x, dnext_h)\n","dprev_h_num = compute_numeric_gradient(fh, h, dnext_h)\n","dWx_num = compute_numeric_gradient(fWx, Wx, dnext_h)\n","dWh_num = compute_numeric_gradient(fWh, Wh, dnext_h)\n","db_num = compute_numeric_gradient(fb, b, dnext_h)\n","\n","dx, dprev_h, dWx, dWh, db = rnn_step_backward(dnext_h, cache)\n","\n","print('dx error: ', rel_error(dx_num, dx))\n","print('dprev_h error: ', rel_error(dprev_h_num, dprev_h))\n","print('dWx error: ', rel_error(dWx_num, dWx))\n","print('dWh error: ', rel_error(dWh_num, dWh))\n","print('db error: ', rel_error(db_num, db))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"vjZjH5JW6bhN"},"source":["## Vanilla RNN: forward\n","Now that you have implemented the forward and backward passes for a single timestep of a vanilla RNN, you will combine these pieces to implement a RNN that processes an entire sequence of data. First implement the forward pass by making calls to the `rnn_step_forward` function that you defined earlier."]},{"cell_type":"code","metadata":{"id":"rOKYPN5BhtJg"},"source":["def rnn_forward(x, h0, Wx, Wh, b):\n","    \"\"\"\n","    Run a vanilla RNN forward on an entire sequence of data. We assume an input\n","    sequence composed of T vectors, each of dimension D. The RNN uses a hidden\n","    size of H, and we work over a minibatch containing N sequences. After running\n","    the RNN forward, we return the hidden states for all timesteps.\n","\n","    Inputs:\n","    - x: Input data for the entire timeseries, of shape (N, T, D).\n","    - h0: Initial hidden state, of shape (N, H)\n","    - Wx: Weight matrix for input-to-hidden connections, of shape (D, H)\n","    - Wh: Weight matrix for hidden-to-hidden connections, of shape (H, H)\n","    - b: Biases, of shape (H,)\n","\n","    Returns a tuple of:\n","    - h: Hidden states for the entire timeseries, of shape (N, T, H).\n","    - cache: Values needed in the backward pass\n","    \"\"\"\n","    h, cache = None, None\n","    ##############################################################################\n","    # TODO: Implement forward pass for a vanilla RNN running on a sequence of    #\n","    # input data. You should use the rnn_step_forward function that you defined  #\n","    # above. You can use a for loop to help compute the forward pass.            #\n","    ##############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","    return h, cache"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"C4VuA3KVh4Tl"},"source":["Run the following to check your implementation. You should see errors on the order of `1e-6` or less."]},{"cell_type":"code","metadata":{"id":"_GQWEn3Z6bhO"},"source":["N, T, D, H = 2, 3, 4, 5\n","\n","x = torch.linspace(-0.1, 0.3, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n","h0 = torch.linspace(-0.3, 0.1, steps=N*H, **to_double_cuda).reshape(N, H)\n","Wx = torch.linspace(-0.2, 0.4, steps=D*H, **to_double_cuda).reshape(D, H)\n","Wh = torch.linspace(-0.4, 0.1, steps=H*H, **to_double_cuda).reshape(H, H)\n","b = torch.linspace(-0.7, 0.1, steps=H, **to_double_cuda)\n","\n","h, _ = rnn_forward(x, h0, Wx, Wh, b)\n","expected_h = torch.tensor([\n","  [\n","    [-0.42070749, -0.27279261, -0.11074945,  0.05740409,  0.22236251],\n","    [-0.39525808, -0.22554661, -0.0409454,   0.14649412,  0.32397316],\n","    [-0.42305111, -0.24223728, -0.04287027,  0.15997045,  0.35014525],\n","  ],\n","  [\n","    [-0.55857474, -0.39065825, -0.19198182,  0.02378408,  0.23735671],\n","    [-0.27150199, -0.07088804,  0.13562939,  0.33099728,  0.50158768],\n","    [-0.51014825, -0.30524429, -0.06755202,  0.17806392,  0.40333043]]], **to_double_cuda)\n","print('h error: ', rel_error(expected_h, h))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"P570PTsw6bhP"},"source":["## Vanilla RNN: backward\n","Implement the backward pass for a vanilla RNN in the function `rnn_backward`. This should run back-propagation over the entire sequence, making calls to the `rnn_step_backward` function that you defined earlier."]},{"cell_type":"code","metadata":{"id":"GESGr4gUh9yP"},"source":["def rnn_backward(dh, cache):\n","    \"\"\"\n","    Compute the backward pass for a vanilla RNN over an entire sequence of data.\n","\n","    Inputs:\n","    - dh: Upstream gradients of all hidden states, of shape (N, T, H). \n","    \n","    NOTE: 'dh' contains the upstream gradients produced by the \n","    individual loss functions at each timestep, *not* the gradients\n","    being passed between timesteps (which you'll have to compute yourself\n","    by calling rnn_step_backward in a loop).\n","\n","    Returns a tuple of:\n","    - dx: Gradient of inputs, of shape (N, T, D)\n","    - dh0: Gradient of initial hidden state, of shape (N, H)\n","    - dWx: Gradient of input-to-hidden weights, of shape (D, H)\n","    - dWh: Gradient of hidden-to-hidden weights, of shape (H, H)\n","    - db: Gradient of biases, of shape (H,)\n","    \"\"\"\n","    dx, dh0, dWx, dWh, db = None, None, None, None, None\n","    ##############################################################################\n","    # TODO: Implement the backward pass for a vanilla RNN running an entire      #\n","    # sequence of data. You should use the rnn_step_backward function that you   #\n","    # defined above. You can use a for loop to help compute the backward pass.   #\n","    ##############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","    return dx, dh0, dWx, dWh, db"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"EX6--VsGiDwK"},"source":[" You should see errors on the order of 1e-7 or less."]},{"cell_type":"code","metadata":{"id":"Ny25RusA6bhQ"},"source":["fix_random_seed(0)\n","\n","N, D, T, H = 2, 3, 10, 5\n","\n","x = torch.randn(N, T, D, **to_double_cuda)\n","h0 = torch.randn(N, H, **to_double_cuda)\n","Wx = torch.randn(D, H, **to_double_cuda)\n","Wh = torch.randn(H, H, **to_double_cuda)\n","b = torch.randn(H, **to_double_cuda)\n","\n","out, cache = rnn_forward(x, h0, Wx, Wh, b)\n","\n","dout = torch.randn(*out.shape, **to_double_cuda)\n","\n","dx, dh0, dWx, dWh, db = rnn_backward(dout, cache)\n","\n","fx = lambda x: rnn_forward(x, h0, Wx, Wh, b)[0]\n","fh0 = lambda h0: rnn_forward(x, h0, Wx, Wh, b)[0]\n","fWx = lambda Wx: rnn_forward(x, h0, Wx, Wh, b)[0]\n","fWh = lambda Wh: rnn_forward(x, h0, Wx, Wh, b)[0]\n","fb = lambda b: rnn_forward(x, h0, Wx, Wh, b)[0]\n","\n","dx_num = compute_numeric_gradient(fx, x, dout)\n","dh0_num = compute_numeric_gradient(fh0, h0, dout)\n","dWx_num = compute_numeric_gradient(fWx, Wx, dout)\n","dWh_num = compute_numeric_gradient(fWh, Wh, dout)\n","db_num = compute_numeric_gradient(fb, b, dout)\n","\n","print('dx error: ', rel_error(dx_num, dx))\n","print('dh0 error: ', rel_error(dh0_num, dh0))\n","print('dWx error: ', rel_error(dWx_num, dWx))\n","print('dWh error: ', rel_error(dWh_num, dWh))\n","print('db error: ', rel_error(db_num, db))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"oEDUmZOkU_LO"},"source":["## Vanilla RNN: backward with autograd\n","Now it's time to introduce the lifesaver PyTorch Automatic Differantiation package - `torch.autograd`!\n","\n","`torch.autograd` provides classes and functions implementing **automatic differentiation** of arbitrary scalar valued functions. It requires minimal changes to the existing code - if you pass tensors with `requires_grad=True` to the forward function you wrote earlier, you can just call `.backward(gradient=grad)` on the output to compute gradients on the input and weights.\n","\n","Now we can compare the manual backward pass with the autograd backward pass. **Read through the following.**  You should get a relative error less than 1e-12."]},{"cell_type":"code","metadata":{"id":"5AMXoqNOVRa_"},"source":["fix_random_seed(0)\n","\n","N, D, T, H = 2, 3, 10, 5\n","\n","# set requires_grad=True\n","x = torch.randn(N, T, D, **to_double_cuda, requires_grad=True)\n","h0 = torch.randn(N, H, **to_double_cuda, requires_grad=True)\n","Wx = torch.randn(D, H, **to_double_cuda, requires_grad=True)\n","Wh = torch.randn(H, H, **to_double_cuda, requires_grad=True)\n","b = torch.randn(H, **to_double_cuda, requires_grad=True)\n","\n","out, cache = rnn_forward(x, h0, Wx, Wh, b)\n","\n","dout = torch.randn(*out.shape, **to_double_cuda)\n","\n","# manual backward\n","with torch.no_grad():\n","  dx, dh0, dWx, dWh, db = rnn_backward(dout, cache)\n","\n","# backward with autograd\n","out.backward(dout) # the magic happens here!\n","dx_auto, dh0_auto, dWx_auto, dWh_auto, db_auto = \\\n","  x.grad, h0.grad, Wx.grad, Wh.grad, b.grad\n","\n","print('dx error: ', rel_error(dx_auto, dx))\n","print('dh0 error: ', rel_error(dh0_auto, dh0))\n","print('dWx error: ', rel_error(dWx_auto, dWx))\n","print('dWh error: ', rel_error(dWh_auto, dWh))\n","print('db error: ', rel_error(db_auto, db))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"zgmxOjX0prA4"},"source":["## RNN Module\n","We can now wrap the vanilla RNN we wrote into an `nn.Module`. `nn.Module` is a base class for all neural network modules, more details regarding its attributes, functions, and methods could be found [here](https://pytorch.org/docs/stable/nn.html?highlight=module#torch.nn.Module).\n","\n","Here we want to set up a module for RNN, where function `__init__` sets up weight and biases, and function `forward` call the `rnn_forward` function from before.\n","\n","**We have written this part for you but you are highly recommended to go through the code as you will write `modules` on your own later.**\n","\n","All the implementation will be with `autograd` and `nn.Module` going forward."]},{"cell_type":"code","metadata":{"id":"8HWcrAiHn2TI"},"source":["class RNN(nn.Module):\n","  \"\"\"\n","  A single-layer vanilla RNN module.\n","  \n","  Arguments for initialization:\n","  - input_size: Input size, denoted as D before\n","  - hidden_size: Hidden size, denoted as H before\n","  \"\"\"\n","  def __init__(self, input_size, hidden_size, device='cpu',\n","                dtype=torch.float32):\n","    \"\"\"\n","    Initialize a RNN.\n","    Model parameters to initialize:\n","    - Wx: Weight matrix for input-to-hidden connections, of shape (D, H)\n","    - Wh: Weight matrix for hidden-to-hidden connections, of shape (H, H)\n","    - b: Biases, of shape (H,)\n","    \"\"\"\n","    super().__init__()\n","    \n","    # Register parameters\n","    self.Wx = Parameter(torch.randn(input_size, hidden_size,\n","                       device=device, dtype=dtype).div(math.sqrt(input_size)))\n","    self.Wh = Parameter(torch.randn(hidden_size, hidden_size,\n","                       device=device, dtype=dtype).div(math.sqrt(hidden_size)))\n","    self.b = Parameter(torch.zeros(hidden_size,\n","                       device=device, dtype=dtype))\n","    \n","  def forward(self, x, h0):\n","    \"\"\"\n","    Inputs:\n","    - x: Input data for the entire timeseries, of shape (N, T, D)\n","    - h0: Initial hidden state, of shape (N, H)\n","\n","    Outputs:\n","    - hn: The hidden state output\n","    \"\"\"\n","    hn, _ = rnn_forward(x, h0, self.Wx, self.Wh, self.b)\n","    return hn\n","  \n","  def step_forward(self, x, prev_h):\n","    \"\"\"\n","    Inputs:\n","    - x: Input data for one time step, of shape (N, D)\n","    - prev_h: The previous hidden state, of shape (N, H)\n","\n","    Outputs:\n","    - next_h: The next hidden state, of shape (N, H)\n","    \"\"\"\n","    next_h, _ = rnn_step_forward(x, prev_h, self.Wx, self.Wh, self.b)\n","    return next_h"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"CIjmnjRd6bhZ"},"source":["# RNN for image captioning\n","You will implement a few necessary tools and layers in order to build an image captioning model (class `CaptioningRNN`)."]},{"cell_type":"markdown","metadata":{"id":"IAa1Kvl2P_2k"},"source":["## Image Feature Extraction\n","Here, we use [MobileNet v2](https://pytorch.org/hub/pytorch_vision_mobilenet_v2/) for image feature extraction. For vanilla RNN and LSTM, we use the pooled CNN feature activation. For Attention LSTM, we use the CNN feature activation map after the last convolution layer."]},{"cell_type":"code","metadata":{"id":"VLmU_CiURha7"},"source":["# !pip install torchsummary\n","\n","class FeatureExtractor(object):\n","  \"\"\"\n","  Image feature extraction with MobileNet.\n","  \"\"\"\n","  def __init__(self, pooling=False, verbose=False,\n","               device='cpu', dtype=torch.float32):\n","\n","    from torchvision import transforms, models\n","    from torchsummary import summary\n","    self.preprocess = transforms.Compose([\n","        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n","    ])\n","    self.device, self.dtype = device, dtype\n","    self.mobilenet = models.mobilenet_v2(pretrained=True).to(device)\n","    self.mobilenet = nn.Sequential(*list(self.mobilenet.children())[:-1]) # Remove the last classifier\n","    \n","    # average pooling\n","    if pooling:\n","      self.mobilenet.add_module('LastAvgPool', nn.AvgPool2d(4, 4)) # input: N x 1280 x 4 x 4\n","    \n","    self.mobilenet.eval()\n","    if verbose:\n","      summary(self.mobilenet, (3, 112, 112))\n","  \n","  def extract_mobilenet_feature(self, img, verbose=False):\n","    \"\"\"\n","    Inputs:\n","    - img: Batch of resized images, of shape N x 3 x 112 x 112\n","\n","    Outputs:\n","    - feat: Image feature, of shape N x 1280 (pooled) or N x 1280 x 4 x 4\n","    \"\"\"\n","    num_img = img.shape[0]\n","    \n","    img_prepro = []\n","    for i in range(num_img):\n","      img_prepro.append(self.preprocess(img[i].type(self.dtype).div(255.)))\n","    img_prepro = torch.stack(img_prepro).to(self.device)\n","    \n","    with torch.no_grad():\n","      feat = []\n","      process_batch = 500\n","      for b in range(math.ceil(num_img/process_batch)):\n","        feat.append(self.mobilenet(img_prepro[b*process_batch:(b+1)*process_batch]\n","                                ).squeeze(-1).squeeze(-1)) # forward and squeeze\n","      feat = torch.cat(feat)\n","      \n","      # add l2 normalization\n","      F.normalize(feat, p=2, dim=1)\n","    \n","    if verbose:\n","      print('Output feature shape: ', feat.shape)\n","    \n","    return feat"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"qHeRMCcjx6v0"},"source":["Now, let's see what's inside MobileNet v2. Assume we have a 3x112x112 image input. We pass argument `pooling=True` to the model so the CNN activation is spatially-pooled from `1280x4x4` to `1280`."]},{"cell_type":"code","metadata":{"id":"_pV0Lau_yDwX"},"source":["model = FeatureExtractor(pooling=True, verbose=True, device='cuda')"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"SVAxU-jO6bhR"},"source":["## Word embedding\n","In deep learning systems, we commonly represent words using vectors. Each word of the vocabulary will be associated with a vector, and these vectors will be learned jointly with the rest of the system.\n","\n","Implement the module `WordEmbedding` to convert words (represented by integers) into vectors."]},{"cell_type":"code","metadata":{"id":"4lIaeSAxiJlY"},"source":["class WordEmbedding(nn.Module):\n","  \"\"\"\n","  Simplified version of torch.nn.Embedding.\n","\n","  We operate on minibatches of size N where\n","  each sequence has length T. We assume a vocabulary of V words, assigning each\n","  word to a vector of dimension D.\n","\n","  Inputs:\n","  - x: Integer array of shape (N, T) giving indices of words. Each element idx\n","    of x muxt be in the range 0 <= idx < V.\n","\n","  Returns a tuple of:\n","  - out: Array of shape (N, T, D) giving word vectors for all input words.\n","  \"\"\"\n","  def __init__(self, vocab_size, embed_size,\n","               device='cpu', dtype=torch.float32):\n","      super().__init__()\n","      \n","      # Register parameters\n","      self.W_embed = Parameter(torch.randn(vocab_size, embed_size,\n","                         device=device, dtype=dtype).div(math.sqrt(vocab_size)))\n","      \n","  def forward(self, x):\n","\n","      out = None\n","      ##############################################################################\n","      # TODO: Implement the forward pass for word embeddings.                      #\n","      #                                                                            #\n","      # HINT: This can be done in one line using PyTorch's array indexing.           #\n","      ##############################################################################\n","      # Replace \"pass\" statement with your code\n","      pass\n","      ##############################################################################\n","      #                               END OF YOUR CODE                             #\n","      ##############################################################################\n","      return out"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"7itJoPZOiL7r"},"source":["Run the following to check your implementation. You should see an error on the order of `1e-7` or less."]},{"cell_type":"code","metadata":{"id":"BZuz2ieE6bhR"},"source":["N, T, V, D = 2, 4, 5, 3\n","\n","x = torch.tensor([[0, 3, 1, 2], [2, 1, 0, 3]], **to_long_cuda)\n","W = torch.linspace(0, 1, steps=V*D, **to_double_cuda).reshape(V, D)\n","\n","model_emb = WordEmbedding(V, D, **to_double_cuda)\n","model_emb.W_embed.data.copy_(W)\n","out = model_emb(x)\n","expected_out = torch.tensor([\n"," [[ 0.,          0.07142857,  0.14285714],\n","  [ 0.64285714,  0.71428571,  0.78571429],\n","  [ 0.21428571,  0.28571429,  0.35714286],\n","  [ 0.42857143,  0.5,         0.57142857]],\n"," [[ 0.42857143,  0.5,         0.57142857],\n","  [ 0.21428571,  0.28571429,  0.35714286],\n","  [ 0.,          0.07142857,  0.14285714],\n","  [ 0.64285714,  0.71428571,  0.78571429]]], **to_double_cuda)\n","\n","print('out error: ', rel_error(expected_out, out))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"6u-4Sr--6bhV","tags":[]},"source":["## (Temporal) Affine layer\n","At every timestep we use an affine function to transform the RNN hidden vector at that timestep into scores for each word in the vocabulary. This could be easily done with the [`nn.Linear`](https://pytorch.org/docs/master/nn.html#torch.nn.Linear) module. It can also work as a regular affine layer, like the one you have implemented from previous assignments. Run the following examples to see how it works. You will intensively use `nn.Linear` later."]},{"cell_type":"code","metadata":{"id":"SA2DgU4X6bhV","tags":[]},"source":["fix_random_seed(0)\n","\n","N, T, D, M = 2, 3, 4, 3\n","\n","w = torch.linspace(-0.2, 0.4, steps=D*M, **to_double_cuda).reshape(D, M).permute(1, 0)\n","b = torch.linspace(-0.4, 0.1, steps=M, **to_double_cuda)\n","\n","temporal_affine = nn.Linear(D, M).to(**to_double_cuda)\n","temporal_affine.weight.data.copy_(w)\n","temporal_affine.bias.data.copy_(b)\n","\n","# For regular affine layer\n","x = torch.linspace(-0.1, 0.3, steps=N*D, **to_double_cuda).reshape(N, D)\n","out = temporal_affine(x)\n","print('affine layer - input shape: {}, output shape: {}'.format(x.shape, out.shape))\n","correct_out = torch.tensor([[-0.35584416, -0.10896104,  0.13792208],\n","                     [-0.31428571, -0.01753247,  0.27922078]], **to_double_cuda)\n","\n","print('dx error: ', rel_error(out, correct_out))\n","\n","\n","# For temporal affine layer\n","x = torch.linspace(-0.1, 0.3, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n","out = temporal_affine(x)\n","print('\\ntemporal affine layer - input shape: {}, output shape: {}'.format(x.shape, out.shape))\n","correct_out = torch.tensor([[[-0.39920949, -0.16533597,  0.06853755],\n","                             [-0.38656126, -0.13750988,  0.11154150],\n","                             [-0.37391304, -0.10968379,  0.15454545]],\n","                            [[-0.36126482, -0.08185771,  0.19754941],\n","                             [-0.34861660, -0.05403162,  0.24055336],\n","                             [-0.33596838, -0.02620553,  0.28355731]]], **to_double_cuda)\n","\n","print('dx error: ', rel_error(out, correct_out))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"K6Py13Ak6bhX","tags":[]},"source":["## Temporal Softmax loss\n","In an RNN language model, at every timestep we produce a score for each word in the vocabulary. We know the ground-truth word at each timestep, so we use a softmax loss function to compute loss and gradient at each timestep. We sum the losses over time and average them over the minibatch.\n","\n","However there is one wrinkle: since we operate over minibatches and different captions may have different lengths, we append `<NULL>` tokens to the end of each caption so they all have the same length. We don't want these `<NULL>` tokens to count toward the loss or gradient, so in addition to scores and ground-truth labels our loss function also accepts a `ignore_index` that tells it which index in caption should be ignored when computing the loss."]},{"cell_type":"code","metadata":{"id":"JyMWheteYQDh"},"source":["def temporal_softmax_loss(x, y, ignore_index=NULL_index):\n","    \"\"\"\n","    A temporal version of softmax loss for use in RNNs. We assume that we are\n","    making predictions over a vocabulary of size V for each timestep of a\n","    timeseries of length T, over a minibatch of size N. The input x gives scores\n","    for all vocabulary elements at all timesteps, and y gives the indices of the\n","    ground-truth element at each timestep. We use a cross-entropy loss at each\n","    timestep, *summing* the loss over all timesteps and *averaging* across the\n","    minibatch.\n","\n","    As an additional complication, we may want to ignore the model output at some\n","    timesteps, since sequences of different length may have been combined into a\n","    minibatch and padded with NULL tokens. The optional ignore_index argument\n","    tells us which elements in the caption should not contribute to the loss.\n","\n","    Inputs:\n","    - x: Input scores, of shape (N, T, V)\n","    - y: Ground-truth indices, of shape (N, T) where each element is in the range\n","         0 <= y[i, t] < V\n","\n","    Returns a tuple of:\n","    - loss: Scalar giving loss\n","    \"\"\"\n","    loss = None\n","    \n","    ##############################################################################\n","    # TODO: Implement the temporal softmax loss function.                        #\n","    #                                                                            #\n","    # REQUIREMENT: This part MUST be done in one single line of code!            #\n","    #                                                                            #\n","    # HINT: Look up the function torch.functional.cross_entropy, set             #\n","    # ignore_index to the variable ignore_index (i.e., index of NULL) and        #\n","    # set reduction to either 'sum' or 'mean' (avoid using 'none' for now).      #\n","    #                                                                            #\n","    # We use a cross-entropy loss at each timestep, *summing* the loss over      #\n","    # all timesteps and *averaging* across the minibatch.                        #\n","    ##############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","    \n","    return loss"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"NmU62YVdLFR0"},"source":["### Sanity check"]},{"cell_type":"code","metadata":{"id":"nlFvgXtD6bhX","tags":[]},"source":["def check_loss(N, T, V, p):\n","    x = 0.001 * torch.randn(N, T, V, **to_double_cuda)\n","    y = torch.randint(V, size=(N, T), **to_long_cuda)\n","    mask = torch.rand(N, T, **to_double_cuda)\n","    y[mask > p] = 0\n","    print(temporal_softmax_loss(x, y).item())\n","  \n","check_loss(1000, 1, 10, 1.0)   # Should be about 2.00-2.11\n","check_loss(1000, 10, 10, 1.0)  # Should be about 20.6-21.0\n","check_loss(5000, 10, 10, 0.1) # Should be about 2.00-2.11"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"XWrmaSZaUxqX"},"source":["## Captioning Module\n","Now we are wrapping everything into the captioning module. Implement the `__init__` function for initialization and the `captioning_forward` for the forward pass. For now you only need to implement for the case where `cell_type='rnn'`, indicating vanialla RNNs; you will implement the `LSTM` case and `AttentionLSTM` case later."]},{"cell_type":"code","metadata":{"id":"1NytlYLekOlJ"},"source":["class CaptioningRNN(nn.Module):\n","    \"\"\"\n","    A CaptioningRNN produces captions from images using a recurrent\n","    neural network.\n","\n","    The RNN receives input vectors of size D, has a vocab size of V, works on\n","    sequences of length T, has an RNN hidden dimension of H, uses word vectors\n","    of dimension W, and operates on minibatches of size N.\n","\n","    Note that we don't use any regularization for the CaptioningRNN.\n","    \n","    You will implement the `__init__` method for model initialization and\n","    the `forward` method first, then come back for the `sample` method later.\n","    \"\"\"\n","    def __init__(self, word_to_idx, input_dim=512, wordvec_dim=128,\n","                 hidden_dim=128, cell_type='rnn', device='cpu', dtype=torch.float32):\n","        \"\"\"\n","        Construct a new CaptioningRNN instance.\n","\n","        Inputs:\n","        - word_to_idx: A dictionary giving the vocabulary. It contains V entries,\n","          and maps each string to a unique integer in the range [0, V).\n","        - input_dim: Dimension D of input image feature vectors.\n","        - wordvec_dim: Dimension W of word vectors.\n","        - hidden_dim: Dimension H for the hidden state of the RNN.\n","        - cell_type: What type of RNN to use; either 'rnn' or 'lstm'.\n","        - dtype: datatype to use; use float32 for training and float64 for\n","          numeric gradient checking.\n","        \"\"\"\n","        super().__init__()\n","        if cell_type not in {'rnn', 'lstm', 'attention'}:\n","            raise ValueError('Invalid cell_type \"%s\"' % cell_type)\n","\n","        self.cell_type = cell_type\n","        self.word_to_idx = word_to_idx\n","        self.idx_to_word = {i: w for w, i in word_to_idx.items()}\n","\n","        vocab_size = len(word_to_idx)\n","\n","        self._null = word_to_idx['<NULL>']\n","        self._start = word_to_idx.get('<START>', None)\n","        self._end = word_to_idx.get('<END>', None)\n","        \n","        ##########################################################################\n","        # TODO: Initialize the image captioning module. Refer to the TODO        #\n","        # in the captioning_forward function on layers you need to create        #\n","        #                                                                        #\n","        # Hint: You can use nn.Linear for both                                   #\n","        # i) output projection (from RNN hidden state to vocab probability) and  #\n","        # ii) feature projection (from CNN pooled feature to h0)                 #\n","        #                                                                        #\n","        # Hint: In FeatureExtractor, set pooling=True to get the pooled CNN      #\n","        #       feature and pooling=False to get the CNN activation map.         #\n","        ##########################################################################\n","        # Replace \"pass\" statement with your code\n","        pass\n","        #############################################################################\n","        #                              END OF YOUR CODE                             #\n","        #############################################################################\n","    \n","    def forward(self):\n","      raise NotImplementedError\n","    \n","    def sample(self):\n","      raise NotImplementedError"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"e2AJE8HVELe-"},"source":["### Forward part\n","Implement the forward function."]},{"cell_type":"code","metadata":{"id":"0my9LmoD0bEh"},"source":["    def captioning_forward(self, images, captions):\n","        \"\"\"\n","        Compute training-time loss for the RNN. We input images and\n","        ground-truth captions for those images, and use an RNN (or LSTM) to compute\n","        loss. The backward part will be done by torch.autograd.\n","\n","        Inputs:\n","        - images: Input images, of shape (N, 3, 112, 112)\n","        - captions: Ground-truth captions; an integer array of shape (N, T + 1) where\n","          each element is in the range 0 <= y[i, t] < V\n","\n","        Outputs:\n","        - loss: A scalar loss\n","        \"\"\"\n","        # Cut captions into two pieces: captions_in has everything but the last word\n","        # and will be input to the RNN; captions_out has everything but the first\n","        # word and this is what we will expect the RNN to generate. These are offset\n","        # by one relative to each other because the RNN should produce word (t+1)\n","        # after receiving word t. The first element of captions_in will be the START\n","        # token, and the first element of captions_out will be the first word.\n","        captions_in = captions[:, :-1]\n","        captions_out = captions[:, 1:]\n","\n","        loss = 0.0\n","        ############################################################################\n","        # TODO: Implement the forward pass for the CaptioningRNN.                  #\n","        # In the forward pass you will need to do the following:                   #\n","        # (1) Use an affine transformation to project the image feature to         #\n","        #     the initial hidden state $h0$ (for RNN/LSTM, of shape (N, H)) or     #\n","        #     the projected CNN activation input $A$ (for Attention LSTM,          #\n","        #     of shape (N, H, 4, 4).                                               #\n","        # (2) Use a word embedding layer to transform the words in captions_in     #\n","        #     from indices to vectors, giving an array of shape (N, T, W).         #\n","        # (3) Use either a vanilla RNN or LSTM (depending on self.cell_type) to    #\n","        #     process the sequence of input word vectors and produce hidden state  #\n","        #     vectors for all timesteps, producing an array of shape (N, T, H).    #\n","        # (4) Use a (temporal) affine transformation to compute scores over the    #\n","        #     vocabulary at every timestep using the hidden states, giving an      #\n","        #     array of shape (N, T, V).                                            #\n","        # (5) Use (temporal) softmax to compute loss using captions_out, ignoring  #\n","        #     the points where the output word is <NULL>.                          #\n","        #                                                                          #\n","        # Do not worry about regularizing the weights or their gradients!          #\n","        ############################################################################\n","        # Replace \"pass\" statement with your code\n","        pass\n","        ############################################################################\n","        #                             END OF YOUR CODE                             #\n","        ############################################################################\n","\n","        return loss\n","\n","    CaptioningRNN.forward = captioning_forward"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"v0M8Dq5NEOxN"},"source":["### Inference part\n","We will come back to this part later."]},{"cell_type":"code","metadata":{"id":"zAZOTp1hz9db"},"source":["    def sample_caption(self, images, max_length=15):\n","        \"\"\"\n","        Run a test-time forward pass for the model, sampling captions for input\n","        feature vectors.\n","\n","        At each timestep, we embed the current word, pass it and the previous hidden\n","        state to the RNN to get the next hidden state, use the hidden state to get\n","        scores for all vocab words, and choose the word with the highest score as\n","        the next word. The initial hidden state is computed by applying an affine\n","        transform to the image features, and the initial word is the <START>\n","        token.\n","\n","        For LSTMs you will also have to keep track of the cell state; in that case\n","        the initial cell state should be zero.\n","\n","        Inputs:\n","        - images: Input images, of shape (N, 3, 112, 112)\n","        - max_length: Maximum length T of generated captions\n","\n","        Returns:\n","        - captions: Array of shape (N, max_length) giving sampled captions,\n","          where each element is an integer in the range [0, V). The first element\n","          of captions should be the first sampled word, not the <START> token.\n","        \"\"\"\n","        N = images.shape[0]\n","        captions = self._null * images.new(N, max_length).fill_(1).long()\n","\n","        if self.cell_type == 'attention':\n","          attn_weights_all = images.new(N, max_length, 4, 4).fill_(0).float()\n","\n","        ###########################################################################\n","        # TODO: Implement test-time sampling for the model. You will need to      #\n","        # initialize the hidden state of the RNN by applying the learned affine   #\n","        # transform to the image features. The first word that you feed to  #\n","        # the RNN should be the <START> token; its value is stored in the         #\n","        # variable self._start. At each timestep you will need to do to:          #\n","        # (1) Embed the previous word using the learned word embeddings           #\n","        # (2) Make an RNN step using the previous hidden state and the embedded   #\n","        #     current word to get the next hidden state.                          #\n","        # (3) Apply the learned affine transformation to the next hidden state to #\n","        #     get scores for all words in the vocabulary                          #\n","        # (4) Select the word with the highest score as the next word, writing it #\n","        #     (the word index) to the appropriate slot in the captions variable   #\n","        #                                                                         #\n","        # For simplicity, you do not need to stop generating after an <END> token #\n","        # is sampled, but you can if you want to.                                 #\n","        #                                                                         #\n","        # HINT: You will not be able to use the rnn_forward or lstm_forward       #\n","        # functions; you'll need to call the `step_forward` from the              #\n","        # RNN/LSTM/AttentionLSTM module in a loop.                                #\n","        #                                                                         #\n","        # NOTE: we are still working over minibatches in this function. Also if   #\n","        # you are using an LSTM, initialize the first cell state to zeros.        #\n","        # For AttentionLSTM, first project the 1280x4x4 CNN feature activation to #\n","        # $A$ of shape Hx4x4. The LSTM initial hidden state and cell state        #\n","        # would both be A.mean(dim=(2, 3)).                                       #\n","        ###########################################################################\n","        # Replace \"pass\" statement with your code\n","        pass\n","        ############################################################################\n","        #                             END OF YOUR CODE                             #\n","        ############################################################################\n","        if self.cell_type == 'attention':\n","          return captions, attn_weights_all.cpu()\n","        else:\n","          return captions\n","\n","    CaptioningRNN.sample = sample_caption"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"CE3Kg3RHLM4o"},"source":["### Sanity check\n","Run the following to check your forward pass using a small test case; you should see difference on the order of 1e-7 or less."]},{"cell_type":"code","metadata":{"id":"d8a71FL_6bhZ"},"source":["fix_random_seed(0)\n","\n","N, D, W, H = 10, 1280, 30, 40\n","D_img = 112\n","word_to_idx = {'<NULL>': 0, 'cat': 2, 'dog': 3}\n","V = len(word_to_idx)\n","T = 13\n","\n","model = CaptioningRNN(word_to_idx,\n","          input_dim=D,\n","          wordvec_dim=W,\n","          hidden_dim=H,\n","          cell_type='rnn',\n","          **to_float_cuda) # use float here to be consistent with MobileNet v2\n","\n","\n","for k,v in model.named_parameters():\n","  # print(k, v.shape) # uncomment this to see the weight shape\n","  v.data.copy_(torch.linspace(-1.4, 1.3, steps=v.numel()).reshape(*v.shape))\n","\n","images = torch.linspace(-3., 3., steps=(N * 3 * D_img * D_img),\n","                       **to_float_cuda).reshape(N, 3, D_img, D_img)\n","captions = (torch.arange(N * T, **to_long_cuda) % V).reshape(N, T)\n","\n","loss = model(images, captions).item()\n","expected_loss = 150.6090393066\n","\n","print('loss: ', loss)\n","print('expected loss: ', expected_loss)\n","print('difference: ', rel_error(torch.tensor(loss), torch.tensor(expected_loss)))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"7YAOcQ4h6bhc"},"source":["## Image Captioning solver\n","Different from the `Solver` class that we used to train image classification models on the previous assignment, on this assignment we use the `torch.optim` package to train image captioning models.\n","\n","We have written this part for you and you need to train the model and generate plots on the training loss."]},{"cell_type":"code","metadata":{"id":"bRC0lNgOxheT"},"source":["def CaptioningTrain(rnn_model, image_data, caption_data, lr_decay=1, **kwargs):\n","  \"\"\"\n","  Run optimization to train the model.\n","  \"\"\"\n","  # optimizer setup\n","  from torch import optim\n","  optimizer = optim.Adam(\n","    filter(lambda p: p.requires_grad, rnn_model.parameters()),\n","    learning_rate) # leave betas and eps by default\n","  lr_scheduler = optim.lr_scheduler.LambdaLR(optimizer,\n","                                             lambda epoch: lr_decay ** epoch)\n","\n","  # sample minibatch data\n","  iter_per_epoch = math.ceil(image_data.shape[0] // batch_size)\n","  loss_history = []\n","  rnn_model.train()\n","  for i in range(num_epochs):\n","    start_t = time.time()\n","    for j in range(iter_per_epoch):\n","      images, captions = image_data[j*batch_size:(j+1)*batch_size], \\\n","                           caption_data[j*batch_size:(j+1)*batch_size]\n","\n","      loss = rnn_model(images, captions)\n","      optimizer.zero_grad()\n","      loss.backward()\n","      loss_history.append(loss.item())\n","      optimizer.step()\n","    end_t = time.time()\n","    print('(Epoch {} / {}) loss: {:.4f} time per epoch: {:.1f}s'.format(\n","        i, num_epochs, loss.item(), end_t-start_t))\n","\n","    lr_scheduler.step()\n","\n","  # plot the training losses\n","  plt.plot(loss_history)\n","  plt.xlabel('Iteration')\n","  plt.ylabel('Loss')\n","  plt.title('Training loss history')\n","  plt.show()"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"eUc5c217pFYE"},"source":["## Overfit small data\n","Once you have familiarized yourself with the `optim` API used above, run the following to make sure your model overfits a small sample of 50 training examples. You should see a final loss of less than 1.6."]},{"cell_type":"code","metadata":{"id":"yzhsGRzk6bhd"},"source":["fix_random_seed(0)\n","\n","# data input\n","small_num_train = 50\n","sample_idx = torch.linspace(0, num_train-1, steps=small_num_train, **to_float_cuda).long()\n","small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n","small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n","\n","# optimization arguments\n","num_epochs = 80\n","batch_size = 50\n","  \n","# create the image captioning model\n","model = CaptioningRNN(\n","          cell_type='rnn',\n","          word_to_idx=data_dict['vocab']['token_to_idx'],\n","          input_dim=1280, # hard-coded, do not modify\n","          hidden_dim=512,\n","          wordvec_dim=256,\n","          **to_float_cuda)\n","\n","for learning_rate in [1e-3]:\n","  print('learning rate is: ', learning_rate)\n","  CaptioningTrain(model, small_image_data, small_caption_data,\n","                num_epochs=num_epochs, batch_size=batch_size,\n","                learning_rate=learning_rate)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"UiHsRysE6bhe"},"source":["## Caption sampling\n","Unlike classification models, image captioning models behave very differently at training time and at test time. At training time, we have access to the ground-truth caption, so we feed ground-truth words as input to the RNN at each timestep. At test time, we sample from the distribution over the vocabulary at each timestep, and feed the sample as input to the RNN at the next timestep.\n","\n","Implement the [`sample`](#scrollTo=v0M8Dq5NEOxN) method in captioning module `CaptioningRNN` for test-time sampling. After doing so, run the following to train a captioning model and sample from the model on both training and validation data."]},{"cell_type":"markdown","metadata":{"id":"BN_sn1YcFZz4"},"source":["### Train the net\n","After you are done implementing the [`sample`](#scrollTo=v0M8Dq5NEOxN) method, perform the training on the entire training set. You should see a final loss less than `2.0`."]},{"cell_type":"code","metadata":{"id":"dXHnPuM_FU7k"},"source":["fix_random_seed(0)\n","\n","# data input\n","small_num_train = num_train\n","sample_idx = torch.randint(num_train, size=(small_num_train,), **to_long_cuda)\n","small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n","small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n","\n","# optimization arguments\n","num_epochs = 60\n","batch_size = 250\n","\n","# create the image captioning model\n","rnn_model = CaptioningRNN(\n","          cell_type='rnn',\n","          word_to_idx=data_dict['vocab']['token_to_idx'],\n","          input_dim=1280, # hard-coded, do not modify\n","          hidden_dim=512,\n","          wordvec_dim=256,\n","          **to_float_cuda)\n","\n","for learning_rate in [1e-3]:\n","  print('learning rate is: ', learning_rate)\n","  CaptioningTrain(rnn_model, small_image_data, small_caption_data,\n","                num_epochs=num_epochs, batch_size=batch_size,\n","                learning_rate=learning_rate)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"97xga3Q5GO8B"},"source":["### Test-time sampling\n","The samples on training data should be very good; the samples on validation data will probably make less sense."]},{"cell_type":"code","metadata":{"id":"Rvt326nX6bhf"},"source":["# Sample a minibatch and show the reshaped 112x112 images,\n","# GT captions, and generated captions by your model.\n","batch_size = 3\n","\n","for split in ['train', 'val']:\n","  sample_idx = torch.randint(0, num_train if split=='train' else num_val, (batch_size,))\n","  sample_images = data_dict[split+'_images'][sample_idx]\n","  sample_captions = data_dict[split+'_captions'][sample_idx]\n","\n","  gt_captions = decode_captions(sample_captions, data_dict['vocab']['idx_to_token'])\n","  rnn_model.eval()\n","  generated_captions = rnn_model.sample(sample_images)\n","  generated_captions = decode_captions(generated_captions, data_dict['vocab']['idx_to_token'])\n","\n","  for i in range(batch_size):\n","    plt.imshow(sample_images[i].permute(1, 2, 0))\n","    plt.axis('off')\n","    plt.title('%s\\nRNN Generated:%s\\nGT:%s' % (split, generated_captions[i], gt_captions[i]))\n","    plt.show()"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"vVeGWWQTvMMF","tags":["pdf-title"]},"source":["# Image Captioning with LSTMs\n","In the previous exercise you implemented a vanilla RNN and applied it to image captioning. Next, we will implement the LSTM update rule and use it for image captioning."]},{"cell_type":"markdown","metadata":{"id":"_8Zd6FGPvMMa"},"source":["# LSTM\n","If you read recent papers, you'll see that many people use a variant on the vanilla RNN called Long-Short Term Memory (LSTM) RNNs. Vanilla RNNs can be tough to train on long sequences due to vanishing and exploding gradients caused by repeated matrix multiplication. LSTMs solve this problem by replacing the simple update rule of the vanilla RNN with a gating mechanism as follows.\n","\n","Similar to the vanilla RNN, at each timestep we receive an input $x_t\\in\\mathbb{R}^D$ and the previous hidden state $h_{t-1}\\in\\mathbb{R}^H$; the LSTM also maintains an $H$-dimensional *cell state*, so we also receive the previous cell state $c_{t-1}\\in\\mathbb{R}^H$. The learnable parameters of the LSTM are an *input-to-hidden* matrix $W_x\\in\\mathbb{R}^{4H\\times D}$, a *hidden-to-hidden* matrix $W_h\\in\\mathbb{R}^{4H\\times H}$ and a *bias vector* $b\\in\\mathbb{R}^{4H}$.\n","\n","At each timestep we first compute an *activation vector* $a\\in\\mathbb{R}^{4H}$ as $a=W_xx_t + W_hh_{t-1}+b$. We then divide this into four vectors $a_i,a_f,a_o,a_g\\in\\mathbb{R}^H$ where $a_i$ consists of the first $H$ elements of $a$, $a_f$ is the next $H$ elements of $a$, etc. We then compute the *input gate* $g\\in\\mathbb{R}^H$, *forget gate* $f\\in\\mathbb{R}^H$, *output gate* $o\\in\\mathbb{R}^H$ and *block input* $g\\in\\mathbb{R}^H$ as\n","\n","$$\n","\\begin{align*}\n","i = \\sigma(a_i) \\hspace{2pc}\n","f = \\sigma(a_f) \\hspace{2pc}\n","o = \\sigma(a_o) \\hspace{2pc}\n","g = \\tanh(a_g)\n","\\end{align*}\n","$$\n","\n","where $\\sigma$ is the sigmoid function and $\\tanh$ is the hyperbolic tangent, both applied elementwise.\n","\n","Finally we compute the next cell state $c_t$ and next hidden state $h_t$ as\n","\n","$$\n","c_{t} = f\\odot c_{t-1} + i\\odot g \\hspace{4pc}\n","h_t = o\\odot\\tanh(c_t)\n","$$\n","\n","where $\\odot$ is the elementwise product of vectors.\n","\n","In the rest of the notebook we will implement the LSTM update rule and apply it to the image captioning task. \n","\n","In the code, we assume that data is stored in batches so that $X_t \\in \\mathbb{R}^{N\\times D}$, and will work with *transposed* versions of the parameters: $W_x \\in \\mathbb{R}^{D \\times 4H}$, $W_h \\in \\mathbb{R}^{H\\times 4H}$ so that activations $A \\in \\mathbb{R}^{N\\times 4H}$ can be computed efficiently as $A = X_t W_x + H_{t-1} W_h$"]},{"cell_type":"markdown","metadata":{"id":"t4DNkZYevMMc"},"source":["## LSTM: step forward\n","Implement the forward pass for a single timestep of an LSTM in the `lstm_step_forward` function. This should be similar to the `rnn_step_forward` function that you implemented above, but using the LSTM update rule instead.\n","\n","Don't worry about the backward part! `autograd` will handle it."]},{"cell_type":"code","metadata":{"id":"7NvdQad46KOQ"},"source":["def lstm_step_forward(x, prev_h, prev_c, Wx, Wh, b, attn=None, Wattn=None):\n","    \"\"\"\n","    Forward pass for a single timestep of an LSTM.\n","\n","    The input data has dimension D, the hidden state has dimension H, and we use\n","    a minibatch size of N.\n","\n","    Inputs:\n","    - x: Input data, of shape (N, D)\n","    - prev_h: Previous hidden state, of shape (N, H)\n","    - prev_c: previous cell state, of shape (N, H)\n","    - Wx: Input-to-hidden weights, of shape (D, 4H)\n","    - Wh: Hidden-to-hidden weights, of shape (H, 4H)\n","    - b: Biases, of shape (4H,)\n","    - attn and Wattn are for Attention LSTM only, indicate the attention input and\n","      embedding weights for the attention input\n","\n","    Returns a tuple of:\n","    - next_h: Next hidden state, of shape (N, H)\n","    - next_c: Next cell state, of shape (N, H)\n","    \"\"\"\n","    next_h, next_c = None, None\n","    #############################################################################\n","    # TODO: Implement the forward pass for a single timestep of an LSTM.        #\n","    # You may want to use torch.sigmoid() for the sigmoid function.             #\n","    #############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","\n","    return next_h, next_c"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"YyzSMwYt6NsZ"},"source":["Once you are done, run the following to perform a simple test of your implementation. You should see errors on the order of `1e-7` or less."]},{"cell_type":"code","metadata":{"id":"awLF_A5ZvMMd"},"source":["N, D, H = 3, 4, 5\n","x = torch.linspace(-0.4, 1.2, steps=N*D, **to_double_cuda).reshape(N, D)\n","prev_h = torch.linspace(-0.3, 0.7, steps=N*H, **to_double_cuda).reshape(N, H)\n","prev_c = torch.linspace(-0.4, 0.9, steps=N*H, **to_double_cuda).reshape(N, H)\n","Wx = torch.linspace(-2.1, 1.3, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n","Wh = torch.linspace(-0.7, 2.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n","b = torch.linspace(0.3, 0.7, steps=4*H, **to_double_cuda)\n","\n","next_h, next_c = lstm_step_forward(x, prev_h, prev_c, Wx, Wh, b)\n","\n","expected_next_h = torch.tensor([\n","    [ 0.24635157,  0.28610883,  0.32240467,  0.35525807,  0.38474904],\n","    [ 0.49223563,  0.55611431,  0.61507696,  0.66844003,  0.7159181 ],\n","    [ 0.56735664,  0.66310127,  0.74419266,  0.80889665,  0.858299  ]], **to_double_cuda)\n","expected_next_c = torch.tensor([\n","    [ 0.32986176,  0.39145139,  0.451556,    0.51014116,  0.56717407],\n","    [ 0.66382255,  0.76674007,  0.87195994,  0.97902709,  1.08751345],\n","    [ 0.74192008,  0.90592151,  1.07717006,  1.25120233,  1.42395676]], **to_double_cuda)\n","\n","print('next_h error: ', rel_error(expected_next_h, next_h))\n","print('next_c error: ', rel_error(expected_next_c, next_c))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"ErgRQwwzvMMt"},"source":["## LSTM: forward\n","Implement the `lstm_forward` function to run an LSTM forward on an entire timeseries of data.\n","\n","Again, don't worry about the backward part! `autograd` will handle it."]},{"cell_type":"code","metadata":{"id":"mAKcjWZI6T5M"},"source":["def lstm_forward(x, h0, Wx, Wh, b):\n","    \"\"\"\n","    Forward pass for an LSTM over an entire sequence of data. We assume an input\n","    sequence composed of T vectors, each of dimension D. The LSTM uses a hidden\n","    size of H, and we work over a minibatch containing N sequences. After running\n","    the LSTM forward, we return the hidden states for all timesteps.\n","\n","    Note that the initial cell state is passed as input, but the initial cell\n","    state is set to zero. Also note that the cell state is not returned; it is\n","    an internal variable to the LSTM and is not accessed from outside.\n","\n","    Inputs:\n","    - x: Input data, of shape (N, T, D)\n","    - h0: Initial hidden state, of shape (N, H)\n","    - Wx: Weights for input-to-hidden connections, of shape (D, 4H)\n","    - Wh: Weights for hidden-to-hidden connections, of shape (H, 4H)\n","    - b: Biases, of shape (4H,)\n","\n","    Returns a tuple of:\n","    - h: Hidden states for all timesteps of all sequences, of shape (N, T, H)\n","    \"\"\"\n","    h = None\n","    c0 = torch.zeros_like(h0) # we provide the intial cell state c0 here for you!\n","    #############################################################################\n","    # TODO: Implement the forward pass for an LSTM over an entire timeseries.   #\n","    # You should use the lstm_step_forward function that you just defined.      #\n","    #############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","\n","    return h"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"eO8f4KcN6XQJ"},"source":["When you are done, run the following to check your implementation. You should see an error on the order of `1e-7` or less."]},{"cell_type":"code","metadata":{"id":"_x-3BJiEvMMv"},"source":["N, D, H, T = 2, 5, 4, 3\n","x = torch.linspace(-0.4, 0.6, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n","h0 = torch.linspace(-0.4, 0.8, steps=N*H, **to_double_cuda).reshape(N, H)\n","Wx = torch.linspace(-0.2, 0.9, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n","Wh = torch.linspace(-0.3, 0.6, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n","b = torch.linspace(0.2, 0.7, steps=4*H, **to_double_cuda)\n","\n","h = lstm_forward(x, h0, Wx, Wh, b)\n","\n","expected_h = torch.tensor([\n"," [[ 0.01764008,  0.01823233,  0.01882671,  0.0194232 ],\n","  [ 0.11287491,  0.12146228,  0.13018446,  0.13902939],\n","  [ 0.31358768,  0.33338627,  0.35304453,  0.37250975]],\n"," [[ 0.45767879,  0.4761092,   0.4936887,   0.51041945],\n","  [ 0.6704845,   0.69350089,  0.71486014,  0.7346449 ],\n","  [ 0.81733511,  0.83677871,  0.85403753,  0.86935314]]], **to_double_cuda)\n","\n","print('h error: ', rel_error(expected_h, h))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"_YZ3OqrqkjLt"},"source":["## LSTM Module\n","\n","We can now wrap the LSTM functions we wrote into an nn.Module."]},{"cell_type":"code","metadata":{"id":"FFgOHHZnkZ53"},"source":["class LSTM(nn.Module):\n","  \"\"\"\n","  This is our single-layer, uni-directional LSTM module.\n","  \n","  Arguments for initialization:\n","  - input_size: Input size, denoted as D before\n","  - hidden_size: Hidden size, denoted as H before\n","  \"\"\"\n","  def __init__(self, input_size, hidden_size, device='cpu',\n","                dtype=torch.float32):\n","    \"\"\"\n","    Initialize a LSTM.\n","    Model parameters to initialize:\n","    - Wx: Weights for input-to-hidden connections, of shape (D, 4H)\n","    - Wh: Weights for hidden-to-hidden connections, of shape (H, 4H)\n","    - b: Biases, of shape (4H,)\n","    \"\"\"\n","    super().__init__()\n","    \n","    # Register parameters\n","    self.Wx = Parameter(torch.randn(input_size, hidden_size*4,\n","                       device=device, dtype=dtype).div(math.sqrt(input_size)))\n","    self.Wh = Parameter(torch.randn(hidden_size, hidden_size*4,\n","                       device=device, dtype=dtype).div(math.sqrt(hidden_size)))\n","    self.b = Parameter(torch.zeros(hidden_size*4,\n","                       device=device, dtype=dtype))\n","    \n","  def forward(self, x, h0):\n","    \"\"\"\n","    Inputs:\n","    - x: Input data for the entire timeseries, of shape (N, T, D)\n","    - h0: Initial hidden state, of shape (N, H)\n","\n","    Outputs:\n","    - hn: The hidden state output\n","    \"\"\"\n","    hn = lstm_forward(x, h0, self.Wx, self.Wh, self.b)\n","    return hn\n","  \n","  def step_forward(self, x, prev_h, prev_c):\n","    \"\"\"\n","    Inputs:\n","    - x: Input data for one time step, of shape (N, D)\n","    - prev_h: The previous hidden state, of shape (N, H)\n","    - prev_c: The previous cell state, of shape (N, H)\n","\n","    Outputs:\n","    - next_h: The next hidden state, of shape (N, H)\n","    - next_c: The next cell state, of shape (N, H)\n","    \"\"\"\n","    next_h, next_c = lstm_step_forward(x, prev_h, prev_c, self.Wx, self.Wh, self.b)\n","    return next_h, next_c"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"92SStL_tvMNK"},"source":["## LSTM captioning model\n","\n","Now that you have implemented an LSTM, update the implementation of the [`init`](#scrollTo=XWrmaSZaUxqX) method in class `CaptioningRNN` **ONLY** to also handle the case where `self.cell_type` is `lstm`. **This should require adding less than 5 lines of code.**\n","\n","Once you have done so, run the following to check your implementation. You should see a difference on the order of `1e-7` or less."]},{"cell_type":"code","metadata":{"id":"NNpiC4WSvMNL"},"source":["fix_random_seed(0)\n","\n","N, D, W, H = 10, 1280, 30, 40\n","D_img = 112\n","word_to_idx = {'<NULL>': 0, 'cat': 2, 'dog': 3}\n","V = len(word_to_idx)\n","T = 13\n","\n","model = CaptioningRNN(word_to_idx,\n","          input_dim=D,\n","          wordvec_dim=W,\n","          hidden_dim=H,\n","          cell_type='lstm',\n","          **to_float_cuda)\n","\n","for k,v in model.named_parameters():\n","  # print(k, v.shape) # uncomment this to see the weight shape\n","  v.data.copy_(torch.linspace(-1.4, 1.3, steps=v.numel()).reshape(*v.shape))\n","\n","images = torch.linspace(-3., 3., steps=(N * 3 * D_img * D_img),\n","                       **to_float_cuda).reshape(N, 3, D_img, D_img)\n","captions = (torch.arange(N * T, **to_long_cuda) % V).reshape(N, T)\n","\n","loss = model(images, captions).item()\n","expected_loss = 146.3161468505\n","\n","print('loss: ', loss)\n","print('expected loss: ', expected_loss)\n","print('difference: ', rel_error(torch.tensor(loss), torch.tensor(expected_loss)))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"06hbDnRXvMNO"},"source":["## Overfit small data\n","We have written this part for you. Run the following to overfit an LSTM captioning model on the same small dataset as we used for the RNN previously. You should see a final loss less than `4` after 80 epochs."]},{"cell_type":"code","metadata":{"id":"O-tETnd3vMNP"},"source":["fix_random_seed(0)\n","\n","# data input\n","small_num_train = 50\n","sample_idx = torch.linspace(0, num_train-1, steps=small_num_train, **to_float_cuda).long()\n","small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n","small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n","\n","# optimization arguments\n","num_epochs = 80\n","batch_size = 50\n","\n","# create the image captioning model\n","model = CaptioningRNN(\n","          cell_type='lstm',\n","          word_to_idx=data_dict['vocab']['token_to_idx'],\n","          input_dim=1280, # hard-coded, do not modify\n","          hidden_dim=512,\n","          wordvec_dim=256,\n","          **to_float_cuda)\n","\n","for learning_rate in [1e-2]:\n","  print('learning rate is: ', learning_rate)\n","  CaptioningTrain(model, small_image_data, small_caption_data,\n","                num_epochs=num_epochs, batch_size=batch_size,\n","                learning_rate=learning_rate)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"4vzLUzlWvMNT"},"source":["## Caption sampling\n","Modify the [`sample`](#scrollTo=v0M8Dq5NEOxN) method in class `CaptioningRNN` to handle the case where `self.cell_type` is `lstm`. **This should take fewer than 10 lines of code.**\n","\n","When you are done, run the following to train a captioning model and sample from your the model on some training and validation set samples."]},{"cell_type":"markdown","metadata":{"id":"bn2PpJL5oC0J"},"source":["### Train the net\n","Now, perform the training on the entire training set. You should see a final loss less than `2.8`."]},{"cell_type":"code","metadata":{"id":"f9MFRowdoHW7"},"source":["fix_random_seed(0)\n","\n","# data input\n","small_num_train = num_train\n","sample_idx = torch.randint(num_train, size=(small_num_train,), **to_long_cuda)\n","small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n","small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n","\n","# optimization arguments\n","num_epochs = 60\n","batch_size = 250\n","\n","# create the image captioning model\n","lstm_model = CaptioningRNN(\n","          cell_type='lstm',\n","          word_to_idx=data_dict['vocab']['token_to_idx'],\n","          input_dim=1280, # hard-coded, do not modify\n","          hidden_dim=512,\n","          wordvec_dim=256,\n","          **to_float_cuda)\n","\n","for learning_rate in [1e-3]:\n","  print('learning rate is: ', learning_rate)\n","  CaptioningTrain(lstm_model, small_image_data, small_caption_data,\n","                num_epochs=num_epochs, batch_size=batch_size,\n","                learning_rate=learning_rate)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"wsM2pIYpG3v1"},"source":["### Test-time sampling\n","As with the RNN, the samples on training data should be very good; the samples on validation data will probably make less sense."]},{"cell_type":"code","metadata":{"id":"ziQJ7SBnvMNU"},"source":["# Sample a minibatch and show the reshaped 112x112 images,\n","# GT captions, and generated captions by your model.\n","batch_size = 3\n","\n","for split in ['train', 'val']:\n","  sample_idx = torch.randint(0, num_train if split=='train' else num_val, (batch_size,))\n","  sample_images = data_dict[split+'_images'][sample_idx]\n","  sample_captions = data_dict[split+'_captions'][sample_idx]\n","\n","  gt_captions = decode_captions(sample_captions, data_dict['vocab']['idx_to_token'])\n","  lstm_model.eval()\n","  generated_captions = lstm_model.sample(sample_images)\n","  generated_captions = decode_captions(generated_captions, data_dict['vocab']['idx_to_token'])\n","\n","  for i in range(batch_size):\n","    plt.imshow(sample_images[i].permute(1, 2, 0))\n","    plt.axis('off')\n","    plt.title('%s\\nLSTM Generated:%s\\nGT:%s' % (split, generated_captions[i], gt_captions[i]))\n","    plt.show()"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"ECcPPE_Pqc8v"},"source":["# Attention LSTM\n","Attention LSTM essentially adds an attention input $x_{attn}^t\\in\\mathbb{R}^H$ into LSTM, along with $x_t\\in\\mathbb{R}^D$ and the previous hidden state $h_{t-1}\\in\\mathbb{R}^H$.\n","\n","To get the attention input $x_{attn}^t$, here we adopt a method called `scaled dot-product attention`, as covered in the lecture. We first project the CNN feature activation from $\\mathbb{R}^{1280\\times4\\times4}$ to $\\mathbb{R}^{H\\times4\\times4}$ using an affine layer. Given the projected activation $A\\in \\mathbb{R}^{H\\times4\\times4}$ and the LSTM hidden state from the previous time step $h_{t-1}$, we formuate the attention weights on $A$ at time step $t$ as $M_{attn}^t=h_{t-1}A/\\sqrt{H} \\in \\mathbb{R}^{4\\times4}$.\n","\n","To simplify the formuation here, we flatten the spatial dimensions of $A$ and $M_{attn}^t$ which gives $\\tilde{A}\\in \\mathbb{R}^{H\\times16}$ and $\\tilde{M^t}_{attn}=h_{t-1}A\\in \\mathbb{R}^{16}$.\n","We add a **`softmax`** activation function on $\\tilde{M^t}_{attn}$ so that the attention weights at each time step are normalized and sum up to one.\n","\n","The attention embedding given the attention weights is then $x_{attn}^t=\\tilde{A}\\tilde{M^t}_{attn} \\in\\mathbb{R}^H$.\n","\n","You will implement a batch version of the attention layer we have described here."]},{"cell_type":"markdown","metadata":{"id":"GTDk54Q4ubQ1"},"source":["## Scaled dot-product attention\n","Implement the scaled dot-product attention function. Given the LSTM hidden state from the previous time step `prev_h` (or $h_{t-1}$) and the projected CNN feature activation `A`, compute the attention weights `attn_weights` (or $\\tilde{M^t}_{attn}$ with a reshaping to $\\mathbb{R}^{4\\times4}$) attention embedding output `attn` (or $x_{attn}^t$) using the formulation we provided."]},{"cell_type":"code","metadata":{"id":"91JFQb7I6-5b"},"source":["def dot_product_attention(prev_h, A):\n","    \"\"\"\n","    A simple scaled dot-product attention layer.\n","    Inputs:\n","    - prev_h: The LSTM hidden state from the previous time step, of shape (N, H)\n","    - A: **Projected** CNN feature activation, of shape (N, H, 4, 4),\n","         where H is the LSTM hidden state size\n","    \n","    Outputs:\n","    - attn: Attention embedding output, of shape (N, H)\n","    - attn_weights: Attention weights, of shape (N, 4, 4)\n","    \n","    \"\"\"\n","    N, H, D_a, _ = A.shape\n","\n","    attn, attn_weights = None, None\n","    #############################################################################\n","    # TODO: Implement the scaled dot-product attention we described earlier.    #\n","    # You will use this function for `attention_forward` and `sample_caption`   #\n","    # HINT: Make sure you reshape attn_weights back to (N, 4, 4)!               #\n","    #############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","    \n","    return attn, attn_weights"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"JmiZIjU2cDu0"},"source":["When you are done, run the following to check your implementation. You should see an error on the order of `1e-7` or less."]},{"cell_type":"code","metadata":{"id":"irAslXWfaVGw"},"source":["N, H = 2, 5\n","D_a = 4\n","\n","prev_h = torch.linspace(-0.4, 0.6, steps=N*H, **to_double_cuda).reshape(N, H)\n","A = torch.linspace(-0.4, 1.8, steps=N*H*D_a*D_a, **to_double_cuda).reshape(N, H, D_a, D_a)\n","\n","attn, attn_weights = dot_product_attention(prev_h, A)\n","\n","expected_attn = torch.tensor([[-0.29784344, -0.07645979,  0.14492386,  0.36630751,  0.58769115],\n","        [ 0.81412643,  1.03551008,  1.25689373,  1.47827738,  1.69966103]], **to_double_cuda)\n","expected_attn_weights = torch.tensor([[[0.06511126, 0.06475411, 0.06439892, 0.06404568],\n","         [0.06369438, 0.06334500, 0.06299754, 0.06265198],\n","         [0.06230832, 0.06196655, 0.06162665, 0.06128861],\n","         [0.06095243, 0.06061809, 0.06028559, 0.05995491]],\n","\n","        [[0.05717142, 0.05784357, 0.05852362, 0.05921167],\n","         [0.05990781, 0.06061213, 0.06132473, 0.06204571],\n","         [0.06277517, 0.06351320, 0.06425991, 0.06501540],\n","         [0.06577977, 0.06655312, 0.06733557, 0.06812722]]], **to_double_cuda)\n","\n","print('attn error: ', rel_error(expected_attn, attn))\n","print('attn_weights error: ', rel_error(expected_attn_weights, attn_weights))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"2VK_Ixn1qlRo"},"source":["## Attention LSTM: step forward"]},{"cell_type":"markdown","metadata":{"id":"DVddQlj4xwRk"},"source":["Modify the [`lstm_step_forward`](#scrollTo=t4DNkZYevMMc) function from earlier to support the extra attention input `attn` (or $x_{attn}$) and its embedding weight matrix `Wattn` (or $W_{attn}$) in the LSTM cell. Hence, at each timestep the *activation vector* $a\\in\\mathbb{R}^{4H}$ in LSTM cell is formulated as:\n","\n","$a=W_xx_t + W_hh_{t-1}+W_{attn}x_{attn}^t+b$.\n","\n","\n","**This should require adding less than 5 lines of code.**\n","\n","Once you are done, run the following to perform a simple test of your implementation. You should see errors on the order of `1e-8` or less."]},{"cell_type":"code","metadata":{"id":"oaS31Ncf3l0d"},"source":["N, D, H = 3, 4, 5\n","\n","x = torch.linspace(-0.4, 1.2, steps=N*D, **to_double_cuda).reshape(N, D)\n","prev_h = torch.linspace(-0.3, 0.7, steps=N*H, **to_double_cuda).reshape(N, H)\n","prev_c = torch.linspace(-0.4, 0.9, steps=N*H, **to_double_cuda).reshape(N, H)\n","Wx = torch.linspace(-2.1, 1.3, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n","Wh = torch.linspace(-0.7, 2.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n","b = torch.linspace(0.3, 0.7, steps=4*H, **to_double_cuda)\n","attn = torch.linspace(0.6, 1.8, steps=N*H, **to_double_cuda).reshape(N, H)\n","Wattn = torch.linspace(1.3, 4.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n","\n","\n","next_h, next_c = lstm_step_forward(x, prev_h, prev_c, Wx, Wh, b, attn, Wattn)\n","\n","expected_next_h = torch.tensor([\n","    [0.53704256, 0.59980774, 0.65596820, 0.70569729, 0.74932626],\n","    [0.78729857, 0.82010653, 0.84828362, 0.87235677, 0.89283167],\n","    [0.91017981, 0.92483119, 0.93717126, 0.94754073, 0.95623746]], **to_double_cuda)\n","expected_next_c = torch.tensor([\n","    [0.59999328, 0.69285041, 0.78570758, 0.87856479, 0.97142202],\n","    [1.06428558, 1.15714276, 1.24999992, 1.34285708, 1.43571424],\n","    [1.52857143, 1.62142857, 1.71428571, 1.80714286, 1.90000000]], **to_double_cuda)\n","\n","print('next_h error: ', rel_error(expected_next_h, next_h))\n","print('next_c error: ', rel_error(expected_next_c, next_c))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"VS0JHfJ53agv"},"source":["## Attention LSTM: forward\n"]},{"cell_type":"markdown","metadata":{"id":"QRry6hvQ7ywx"},"source":["Now, implement the `attention_forward` function to run an Attention LSTM forward on an entire timeseries of data. You will have to use the `dot_product_attention` function and the `lstm_step_forward` function you implemented.\n","\n","Again, don't worry about the backward part! `autograd` will handle it."]},{"cell_type":"code","metadata":{"id":"z8jxdeCxu5fn"},"source":["def attention_forward(x, A, Wx, Wh, Wattn, b):\n","    \"\"\"\n","    h0 and c0 are same initialized as the global image feature (meanpooled A)\n","    For simplicity, we implement scaled dot-product attention, which means in\n","    Eq. 4 of the paper (https://arxiv.org/pdf/1502.03044.pdf),\n","    f_{att}(a_i, h_{t−1}) equals to the scaled dot product of a_i and h_{t-1}.\n","    \n","    Forward pass for an LSTM over an entire sequence of data. We assume an input\n","    sequence composed of T vectors, each of dimension D. The LSTM uses a hidden\n","    size of H, and we work over a minibatch containing N sequences. After running\n","    the LSTM forward, we return the hidden states for all timesteps.\n","\n","    Note that the initial cell state is passed as input, but the initial cell\n","    state is set to zero. Also note that the cell state is not returned; it is\n","    an internal variable to the LSTM and is not accessed from outside.\n","\n","    Inputs:\n","    - x: Input data, of shape (N, T, D)\n","    - A: **Projected** activation map, of shape (N, H, 4, 4)\n","    - Wx: Weights for input-to-hidden connections, of shape (D, 4H)\n","    - Wh: Weights for hidden-to-hidden connections, of shape (H, 4H)\n","    - Wattn: Weights for attention-to-hidden connections, of shape (H, 4H)\n","    - b: Biases, of shape (4H,)\n","\n","    Returns a tuple of:\n","    - h: Hidden states for all timesteps of all sequences, of shape (N, T, H)\n","    \"\"\"\n","    \n","    h = None\n","    \n","    # The initial hidden state h0 and cell state c0 are initialized differently in\n","    # Attention LSTM from the original LSTM and hence we provided them for you.\n","    h0 = A.mean(dim=(2, 3)) # Initial hidden state, of shape (N, H)\n","    c0 = h0 # Initial cell state, of shape (N, H)\n","\n","    #############################################################################\n","    # TODO: Implement the forward pass for an LSTM over an entire timeseries.   #\n","    # You should use the lstm_step_forward function and dot_product_attention   #\n","    # function that you just defined.                                           #\n","    #############################################################################\n","    # Replace \"pass\" statement with your code\n","    pass\n","    ##############################################################################\n","    #                               END OF YOUR CODE                             #\n","    ##############################################################################\n","\n","    return h"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"5nPR-OYQ4Ohp"},"source":["When you are done, run the following to check your implementation. You should see an error on the order of `1e-8` or less."]},{"cell_type":"code","metadata":{"id":"aB6VU8nl4SmS"},"source":["N, D, H, T = 2, 5, 4, 3\n","D_a = 4\n","\n","x = torch.linspace(-0.4, 0.6, steps=N*T*D, **to_double_cuda).reshape(N, T, D)\n","A = torch.linspace(-0.4, 1.8, steps=N*H*D_a*D_a, **to_double_cuda).reshape(N, H, D_a, D_a)\n","Wx = torch.linspace(-0.2, 0.9, steps=4*D*H, **to_double_cuda).reshape(D, 4 * H)\n","Wh = torch.linspace(-0.3, 0.6, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n","Wattn = torch.linspace(1.3, 4.2, steps=4*H*H, **to_double_cuda).reshape(H, 4 * H)\n","b = torch.linspace(0.2, 0.7, steps=4*H, **to_double_cuda)\n","\n","h = attention_forward(x, A, Wx, Wh, Wattn, b)\n","\n","expected_h = torch.tensor([\n","        [[0.56141729, 0.70274849, 0.80000386, 0.86349400],\n","         [0.89556391, 0.92856726, 0.94950579, 0.96281018],\n","         [0.96792077, 0.97535465, 0.98039623, 0.98392994]],\n","\n","        [[0.95065880, 0.97135490, 0.98344373, 0.99045552],\n","         [0.99317679, 0.99607466, 0.99774317, 0.99870293],\n","         [0.99907382, 0.99946784, 0.99969426, 0.99982435]]], **to_double_cuda)\n","\n","print('h error: ', rel_error(expected_h, h))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"jGdYDBjDqofZ"},"source":["## Attention LSTM Module\n","We can now wrap the Attention LSTM functions we wrote into an nn.Module."]},{"cell_type":"code","metadata":{"id":"x4l3yDPx5gWo"},"source":["class AttentionLSTM(nn.Module):\n","  \"\"\"\n","  This is our single-layer, uni-directional Attention module.\n","  \n","  Arguments for initialization:\n","  - input_size: Input size, denoted as D before\n","  - hidden_size: Hidden size, denoted as H before\n","  \"\"\"\n","  def __init__(self, input_size, hidden_size, device='cpu',\n","                dtype=torch.float32):\n","    \"\"\"\n","    Initialize a LSTM.\n","    Model parameters to initialize:\n","    - Wx: Weights for input-to-hidden connections, of shape (D, 4H)\n","    - Wh: Weights for hidden-to-hidden connections, of shape (H, 4H)\n","    - Wattn: Weights for attention-to-hidden connections, of shape (H, 4H)\n","    - b: Biases, of shape (4H,)\n","    \"\"\"\n","    super().__init__()\n","    \n","    # Register parameters\n","    self.Wx = Parameter(torch.randn(input_size, hidden_size*4,\n","                       device=device, dtype=dtype).div(math.sqrt(input_size)))\n","    self.Wh = Parameter(torch.randn(hidden_size, hidden_size*4,\n","                       device=device, dtype=dtype).div(math.sqrt(hidden_size)))\n","    self.Wattn = Parameter(torch.randn(hidden_size, hidden_size*4,\n","                       device=device, dtype=dtype).div(math.sqrt(hidden_size)))\n","    self.b = Parameter(torch.zeros(hidden_size*4,\n","                       device=device, dtype=dtype))\n","    \n","  def forward(self, x, A):\n","    \"\"\"  \n","    Inputs:\n","    - x: Input data for the entire timeseries, of shape (N, T, D)\n","    - A: The projected CNN feature activation, of shape (N, H, 4, 4)\n","\n","    Outputs:\n","    - hn: The hidden state output\n","    \"\"\"\n","    hn = attention_forward(x, A, self.Wx, self.Wh, self.Wattn, self.b)\n","    return hn\n","  \n","  def step_forward(self, x, prev_h, prev_c, attn):\n","    \"\"\"\n","    Inputs:\n","    - x: Input data for one time step, of shape (N, D)\n","    - prev_h: The previous hidden state, of shape (N, H)\n","    - prev_c: The previous cell state, of shape (N, H)\n","    - attn: The attention embedding, of shape (N, H)\n","\n","    Outputs:\n","    - next_h: The next hidden state, of shape (N, H)\n","    - next_c: The next cell state, of shape (N, H)\n","    \"\"\"\n","    next_h, next_c = lstm_step_forward(x, prev_h, prev_c, self.Wx, self.Wh,\n","                                       self.b, attn=attn, Wattn=self.Wattn)\n","    return next_h, next_c"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"9VzpyHuX6Jzc"},"source":["## Attention LSTM captioning model\n","\n","Now that you have implemented an attention module, update the implementation of the [`init`](#scrollTo=XWrmaSZaUxqX) method and [`forward`](#scrollTo=e2AJE8HVELe-) method in module `CaptioningRNN` to also handle the case where `self.cell_type` is `attention`. **This should require adding less than 10 lines of code.**\n","\n","Once you have done so, run the following to check your implementation. You should see a difference on the order of `1e-7` or less."]},{"cell_type":"code","metadata":{"id":"7VqGqDYw6Jzd"},"source":["fix_random_seed(0)\n","\n","N, D, W, H = 10, 1280, 30, 40\n","D_img = 112\n","word_to_idx = {'<NULL>': 0, 'cat': 2, 'dog': 3}\n","V = len(word_to_idx)\n","T = 13\n","\n","model = CaptioningRNN(word_to_idx,\n","          input_dim=D,\n","          wordvec_dim=W,\n","          hidden_dim=H,\n","          cell_type='attention',\n","          **to_float_cuda)\n","\n","for k,v in model.named_parameters():\n","  # print(k, v.shape) # uncomment this to see the weight shape\n","  v.data.copy_(torch.linspace(-1.4, 1.3, steps=v.numel()).reshape(*v.shape))\n","\n","images = torch.linspace(-3., 3., steps=(N * 3 * D_img * D_img),\n","                       **to_float_cuda).reshape(N, 3, D_img, D_img)\n","captions = (torch.arange(N * T, **to_long_cuda) % V).reshape(N, T)\n","\n","loss = model(images, captions).item()\n","expected_loss = 46.9113769531\n","\n","print('loss: ', loss)\n","print('expected loss: ', expected_loss)\n","print('difference: ', rel_error(torch.tensor(loss), torch.tensor(expected_loss)))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"eYxXTAn4q0wV"},"source":["## Overfit small data\n","We have written this part for you. Run the following to overfit an Attention LSTM captioning model on the same small dataset as we used for the RNN previously. You should see a final loss less than `9`."]},{"cell_type":"code","metadata":{"id":"tlK7lKUgWeDS"},"source":["fix_random_seed(0)\n","\n","# data input\n","small_num_train = 50\n","sample_idx = torch.linspace(0, num_train-1, steps=small_num_train, **to_float_cuda).long()\n","small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n","small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n","\n","# optimization arguments\n","num_epochs = 80\n","batch_size = 50\n","\n","# create the image captioning model\n","model = CaptioningRNN(\n","          cell_type='attention',\n","          word_to_idx=data_dict['vocab']['token_to_idx'],\n","          input_dim=1280, # hard-coded, do not modify\n","          hidden_dim=512,\n","          wordvec_dim=256,\n","          **to_float_cuda)\n","\n","for learning_rate in [1e-3]:\n","  print('learning rate is: ', learning_rate)\n","  CaptioningTrain(model, small_image_data, small_caption_data,\n","                num_epochs=num_epochs, batch_size=batch_size,\n","                learning_rate=learning_rate)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"ircMb7_qq7vB"},"source":["## Caption sampling\n","Modify the [`sample`](#scrollTo=v0M8Dq5NEOxN) method in module `CaptioningRNN` to handle the case where `self.cell_type` is `attention`. **This should take fewer than 10 lines of code.**\n","\n","When you are done run the following to train a captioning model and sample from the model on some training and validation set samples."]},{"cell_type":"markdown","metadata":{"id":"t--pa33Sq4SW"},"source":["### Train the net\n","Now, perform the training on the entire training set. You should see a final loss less than `1.0`."]},{"cell_type":"code","metadata":{"id":"ScBvAfcXdVv4"},"source":["fix_random_seed(0)\n","\n","# data input\n","small_num_train = num_train\n","sample_idx = torch.randint(num_train, size=(small_num_train,), **to_long_cuda)\n","small_image_data = data_dict['train_images'][sample_idx].to('cuda')\n","small_caption_data = data_dict['train_captions'][sample_idx].to('cuda')\n","\n","# optimization arguments\n","num_epochs = 60\n","batch_size = 250\n","\n","# create the image captioning model\n","attn_model = CaptioningRNN(\n","          cell_type='attention',\n","          word_to_idx=data_dict['vocab']['token_to_idx'],\n","          input_dim=1280, # hard-coded, do not modify\n","          hidden_dim=512,\n","          wordvec_dim=256,\n","          **to_float_cuda)\n","\n","for learning_rate in [1e-3]:\n","  print('learning rate is: ', learning_rate)\n","  CaptioningTrain(attn_model, small_image_data, small_caption_data,\n","                num_epochs=num_epochs, batch_size=batch_size,\n","                learning_rate=learning_rate)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"5ham_O1TG_z7"},"source":["### Test-time sampling and visualization\n","As with RNN and LSTM, the samples on training data should be very good; the samples on validation data will probably make less sense.\n","\n","We use the `attention_visualizer` function to visualize the attended regions per generated word. Note that sometimes the attended regions (brighter) might not make much sense particially due to our low resolution image input. In real applications, the attended regions are more accurate."]},{"cell_type":"code","metadata":{"id":"0i8KNWSDSLNu"},"source":["# Sample a minibatch and show the reshaped 112x112 images,\n","# GT captions, and generated captions by your model.\n","\n","batch_size = 3\n","from torchvision.utils import make_grid\n","from torchvision import transforms\n","\n","for split in ['train', 'val']:\n","  sample_idx = torch.randint(0, num_train if split=='train' else num_val, (batch_size,))\n","  sample_images = data_dict[split+'_images'][sample_idx]\n","  sample_captions = data_dict[split+'_captions'][sample_idx]\n","\n","  gt_captions = decode_captions(sample_captions, data_dict['vocab']['idx_to_token'])\n","  attn_model.eval()\n","  generated_captions, attn_weights_all = attn_model.sample(sample_images)\n","  generated_captions = decode_captions(generated_captions, data_dict['vocab']['idx_to_token'])\n","\n","  for i in range(batch_size):\n","    plt.imshow(sample_images[i].permute(1, 2, 0))\n","    plt.axis('off')\n","    plt.title('%s\\nAttention LSTM Generated:%s\\nGT:%s' % (split, generated_captions[i], gt_captions[i]))\n","    plt.show()\n","    \n","    tokens = generated_captions[i].split(' ')\n","    \n","    vis_attn = []\n","    for j in range(len(tokens)):\n","      img = sample_images[i]\n","      attn_weights = attn_weights_all[i][j]\n","      token = tokens[j]\n","      img_copy = attention_visualizer(img, attn_weights, token)\n","      vis_attn.append(transforms.ToTensor()(img_copy))\n","    \n","    plt.rcParams['figure.figsize'] = (20.0, 20.0)\n","    vis_attn = make_grid(vis_attn, nrow=8)\n","    plt.imshow(torch.flip(vis_attn, dims=(0,)).permute(1, 2, 0))\n","    plt.axis('off')\n","    plt.show()\n","    plt.rcParams['figure.figsize'] = (10.0, 8.0)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"ghB8BwfUpmI5"},"source":["## Final checks\n","Make sure you run \"Runtime -> Restart and run all...\" to double check the RNN/LSTM code is still functioning well after all the changes you have made!"]}]}